import numpy as np
from paddleocr import PaddleOCR
import paddle
import cv2
from PIL import Image
import os


class CPUOptimizedOCR:
    def __init__(self, model_dir="./models"):
        # 强制使用CPU
        paddle.set_device('cpu')

        # 初始化OCR引擎
        self.ocr = PaddleOCR(
            use_gpu=False,
            use_angle_cls=True,
            lang='ch',
            enable_mkldnn=True,
            det_model_dir=os.path.join(model_dir, 'ch_PP-OCRv4_det'),
            rec_model_dir=os.path.join(model_dir, 'ch_PP-OCRv4_rec'),
            cls_model_dir=os.path.join(model_dir, 'ch_ppocr_mobile_v2.0_cls'),
            show_log=False
        )

    def recognize(self, file_stream):
        try:
            # 读取图片
            img = Image.open(file_stream).convert('RGB')
            img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

            # 执行OCR
            result = self.ocr.ocr(img, cls=True)

            # 格式化结果
            return [{
                "coordinates": np.array(line[0]).reshape(-1, 2).tolist(),
                "text": line[1][0],
                "confidence": float(line[1][1])
            } for line in result[0] if line]

        except Exception as e:
            return {"error": str(e)}